package com.zhb.zhbaiagent.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 基于文件持久化的对话记忆
 * 序列化：将对象序列化成字节流，写入文件
 * 反序列化：将字节流反序列化成对象
 */
public class FileBasedChatMemory implements ChatMemory {

    private final String BASE_DIR;
    private static final Kryo kryo = new Kryo();

    static {
        // 允许未注册的类
        kryo.setRegistrationRequired(false);
        // 设置实例化策略
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    // 构造对象时，指定文件保存目录
    public FileBasedChatMemory(String dir) {
        this.BASE_DIR = dir;
        File baseDir = new File(dir);
        if (!baseDir.exists()) {
            baseDir.mkdirs();
        }
    }

    // 添加对话，保存到文件
    @Override
    public void add(String conversationId, List<Message> messages) {
        List<Message> conversationMessages = getOrCreateConversation(conversationId);
        conversationMessages.addAll(messages);
        saveConversation(conversationId, conversationMessages);
    }

    // 获取对话，从文件读取
    @Override
    public List<Message> get(String conversationId, int lastN) {
        return getOrCreateConversation(conversationId);
    }

//    @Override
//    public List<Message> get(String conversationId) {
//        return getOrCreateConversation(conversationId);
//    }

    // 清空对话，删除文件
    @Override
    public void clear(String conversationId) {
        File file = getConversationFile(conversationId);
        if (file.exists()) {
            file.delete();
        }
    }

    /**
     * 获取或创建对话对象
     * 如果本地文件存在，则从文件中读取对话信息；否则，创建一个新的对话对象
     *
     * @param conversationId 对话的唯一标识符，用于定位特定的对话文件
     * @return 返回一个装载着消息对象的列表，代表对话的内容
     */
    private List<Message> getOrCreateConversation(String conversationId) {
        // 根据对话ID获取对应的文件对象
        File file = getConversationFile(conversationId);
        // 初始化一个空的消息列表，用于存储对话消息
        List<Message> messages = new ArrayList<>();
        // 检查文件是否存在
        if (file.exists()) {
            try (// 尝试打开文件输入流和Input对象
                 Input input = new Input(new FileInputStream(file))) {
                // 使用Kryo反序列化文件中的数据到消息列表
                messages = kryo.readObject(input, ArrayList.class);
            } catch (IOException e) {
                // 如果发生IO异常，打印异常信息
                e.printStackTrace();
            }
        }
        // 返回消息列表，无论是否从文件中读取到了数据，还是一个空列表
        return messages;
    }

    /**
     * 保存对话到文件
     * 此方法负责将给定的对话消息列表序列化到磁盘上的一个文件中，以便于持久化对话数据
     * 它首先确定保存对话的文件路径，然后使用Kryo序列化库将消息列表写入文件
     *
     * @param conversationId 对话的唯一标识符，用于确定保存对话的文件名
     * @param messages 一个包含多个消息的列表，这些消息属于同一个对话
     */
    private void saveConversation(String conversationId, List<Message> messages) {
        // 根据对话ID获取保存对话的文件实例
        File file = getConversationFile(conversationId);
        try (Output output = new Output(new FileOutputStream(file))) {
            // 使用Kryo序列化工具将消息列表写入文件中
            kryo.writeObject(output, messages);
        } catch (IOException e) {
            // 如果在保存对话过程中发生IO错误，则打印错误堆栈跟踪
            e.printStackTrace();
        }
    }

    /**
     * 根据会话ID获取对应的会话文件对象
     * 该方法用于构造一个指向特定会话文件的File对象它将BASE_DIR作为父目录，
     * 并结合会话ID与".kryo"后缀来形成文件名这种方法便于统一管理和定位会话数据文件
     *
     * @param conversationId 会话的唯一标识符，用于区分不同的会话文件
     * @return File对象，表示特定会话的数据文件
     */
    private File getConversationFile(String conversationId) {
        return new File(BASE_DIR, conversationId + ".kryo");
    }
}